import numpy as np
from elasticsearch import Elasticsearch
from FlagEmbedding import FlagModel
import faiss

device = 0


def build_es_index(records):
    es = Elasticsearch("http://localhost:9200", timeout=3600)
    # print(es.info())
    if not es.indices.exists(index="tcm_dose"):
        es.indices.create(index="tcm_dose")
        for id, record in enumerate(records):
            cure = record["user"]["主治"] if "主治" in record["user"] else None
            constituent = record["user"]["组成"] if "组成" in record["user"] else None
            es.index(index="tcm_dose", id=id, body={"id": id, "主治": cure, "组成": constituent, "record": record})
    return es


def es_retrieve(es_index, match_pair: dict, topk=2):
    """

    :param match_pair: {}
    :return:
    """
    result = []
    es_result = es_index.search(index="tcm_dose", query={"match": match_pair}, size=topk)
    for hit in es_result["hits"]["hits"]:
        result.append((hit["_score"], hit["_source"]["record"]))
    return result


class FaissRAG:
    def __init__(self, records, bge_model, dim):
        self.records = records
        self.dim = dim
        self.bge_model = bge_model
        self.index, self.cure_index, self.constituent_index = self.build_index(records, bge_model, dim)

    @staticmethod
    def build_index(records, bge_model, dim=1024):
        """

        :param vector_base: vectors in a matrix to be indexed
        :return:
        """
        # index = faiss.IndexFlatL2(dim)
        cure_index = faiss.IndexFlatIP(dim)
        constituent_index = faiss.IndexFlatIP(dim)
        index = faiss.IndexFlatIP(dim * 2)
        # print(index.is_trained)
        cure_ids = []
        cure_vecs = []
        constituent_ids = []
        constituent_vecs = []
        total_ids = []
        total_vecs = []
        for id, record in enumerate(records):
            cure = record["主治"]
            constituent = record["组成"]
            if cure:
                cure_vec = bge_model.encode(cure)
                cure_vecs.append(cure_vec)
                cure_ids.append(id)
            if constituent:
                constituent_vec = bge_model.encode(constituent)
                constituent_vecs.append(constituent_vec)
                constituent_ids.append(id)
            if cure and constituent:
                total_vecs.append(np.concatenate([cure_index, constituent_vec], axis=1))
                total_ids.append(id)

        cure_index.add_with_ids(np.array(cure_vecs, dtype=np.float32), cure_ids)
        constituent_index.add_with_ids(np.array(constituent_vec, dtype=np.float32), constituent_ids)
        index.add_with_ids(np.array(total_vecs, dtype=np.float32), total_ids)
        return index, cure_index, constituent_index

    def retrieve(self, xq, topk, retrieve_strategy):
        """
        https://github.com/facebookresearch/faiss/wiki/Getting-started

        :param faiss_index:
        :param xq: query vector
        :param topk: we want to see 4 nearest neighbors
        :return:
        """
        '''
        D, I = faiss_index.search(xb[:5], topk)  # sanity check
        print(I)
        print(D)
        '''
        if retrieve_strategy == 'all':
            assert xq.size(-1) == self.dim * 2
            D, I = self.index.search(xq, topk)
        elif retrieve_strategy == 'cure':
            assert xq.size(-1) == self.dim
            D, I = self.cure_index.search(xq, topk)
        elif retrieve_strategy == 'constituent':
            assert xq.size(-1) == self.dim
            D, I = self.constituent_index.search(xq, topk)
        retrieved_records = [self.records[i] for i in I]
        return I, retrieved_records


def bge_index():
    sentences_1 = ["样例数据-1", "样例数据-2"]
    sentences_2 = ["样例数据-3", "样例数据-4"]
    model = FlagModel('BAAI/bge-large-zh-v1.5',
                      # query_instruction_for_retrieval="为这个句子生成表示以用于检索相关文章：",
                      use_fp16=True)  # Setting use_fp16 to True speeds up computation with a slight performance degradation
    embeddings_1 = model.encode(sentences_1)
    # 对于短query使用
    # q_embeddings = model.encode_queries(queries)
    embeddings_2 = model.encode(sentences_2)
    similarity = embeddings_1 @ embeddings_2.T
    return similarity


if __name__ == "__main__":
    pass
